#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include "memcheck.h"
#include "memcheck/common.h"

static inline uint32_t fetch32(const void *p) { return *(const uint32_t *)p; }
static inline void emit32(uint8_t **buf, uint32_t v) { *(uint32_t *)(*buf) = v; *buf += 4; }

static void emit_movabs_x(uint8_t **out, int reg, uint64_t val)
{
    uint16_t i0 = (uint16_t)(val & 0xFFFF);
    uint16_t i1 = (uint16_t)((val >> 16) & 0xFFFF);
    uint16_t i2 = (uint16_t)((val >> 32) & 0xFFFF);
    uint16_t i3 = (uint16_t)((val >> 48) & 0xFFFF);
    uint32_t movz = 0xD2800000;
    uint32_t movk = 0xF2800000;
    emit32(out, movz | ((uint32_t)i0 << 5) | (reg & 31));
    emit32(out, movk | ((uint32_t)i1 << 5) | (reg & 31) | (1u << 21));
    emit32(out, movk | ((uint32_t)i2 << 5) | (reg & 31) | (2u << 21));
    emit32(out, movk | ((uint32_t)i3 << 5) | (reg & 31) | (3u << 21));
}

static void emit_sub_sp_imm(uint8_t **out, unsigned imm)
{
    uint32_t ins = 0xD1000000 | ((imm & 0xFFF) << 10) | (31u << 5) | 31u;
    emit32(out, ins);
}
static void emit_add_sp_imm(uint8_t **out, unsigned imm)
{
    uint32_t ins = 0x91000000 | ((imm & 0xFFF) << 10) | (31u << 5) | 31u;
    emit32(out, ins);
}
static void emit_str_x_sp_off(uint8_t **out, int reg, unsigned offset_bytes)
{
    unsigned imm12 = (offset_bytes >> 3);
    uint32_t ins = 0xF9000000 | ((imm12 & 0xFFF) << 10) | (31u << 5) | (reg & 31);
    emit32(out, ins);
}
static void emit_ldr_x_sp_off(uint8_t **out, int reg, unsigned offset_bytes)
{
    unsigned imm12 = (offset_bytes >> 3);
    uint32_t ins = 0xF9400000 | ((imm12 & 0xFFF) << 10) | (31u << 5) | (reg & 31);
    emit32(out, ins);
}

static inline void emit_br_reg(uint8_t **out, int reg)
{
    // BR Xreg
    emit32(out, 0xD61F0000 | ((reg & 31) << 5));
}

static int is_ldr_str_unsigned(uint32_t insn, int *is_store, size_t *size)
{
    if (((insn >> 26) & 0x1) != 1) return 0;
    if (((insn >> 27) & 0x1) != 0) return 0;
    if (((insn >> 10) & 0x3) != 0) return 0;
    int opc = (insn >> 22) & 0x3;
    int sz = (insn >> 30) & 0x3;
    if (opc == 0) *is_store = 1; else if (opc == 1) *is_store = 0; else return 0;
    switch (sz) { case 2: *size = 4; break; case 3: *size = 8; break; default: return 0; }
    return 1;
}

static int is_pc_relative_or_branch(uint32_t insn)
{
    // ADRP / ADR
    if ((insn & 0x9F000000u) == 0x90000000u) return 1; // ADRP
    if ((insn & 0x9F000000u) == 0x10000000u) return 1; // ADR
    // LDR (literal)
    if ((insn & 0x3B000000u) == 0x18000000u) return 1; // LDR literal class
    // Unconditional branch (immediate) B / BL
    if ((insn & 0x7C000000u) == 0x14000000u) return 1; // B
    if ((insn & 0xFC000000u) == 0x94000000u) return 1; // BL
    // Conditional branch (immediate)
    if ((insn & 0xFF000010u) == 0x54000000u) return 1; // B.cond
    // Compare and branch (immediate)
    if ((insn & 0x7E000000u) == 0x34000000u) return 1; // CBZ/CBNZ
    // Test and branch (immediate)
    if ((insn & 0x7E000000u) == 0x36000000u) return 1; // TBZ/TBNZ
    return 0;
}

void dbi_log_mem_access(uint64_t pc, uint64_t ea, int is_store, size_t size)
{
    dprintf(2, "[memcheck:a64] pc=0x%lx ea=0x%lx %s size=%zu\n",
            (unsigned long)pc, (unsigned long)ea, is_store?"STORE":"LOAD", size);
}

static void emit_call_log(uint8_t **out, uint64_t pc, uint64_t ea, int is_store, size_t size)
{
    // 保存所有 caller-saved 寄存器，避免被日志调用破坏
    // AAPCS64: x0-x17 为调用者保存（x18 保留），x19-x28 为被调保存
    // 这里保存 x0-x17 和 x30
    emit_sub_sp_imm(out, 160);
    emit_str_x_sp_off(out, 0, 0);
    emit_str_x_sp_off(out, 1, 8);
    emit_str_x_sp_off(out, 2, 16);
    emit_str_x_sp_off(out, 3, 24);
    emit_str_x_sp_off(out, 4, 32);
    emit_str_x_sp_off(out, 5, 40);
    emit_str_x_sp_off(out, 6, 48);
    emit_str_x_sp_off(out, 7, 56);
    emit_str_x_sp_off(out, 8, 64);
    emit_str_x_sp_off(out, 9, 72);
    emit_str_x_sp_off(out, 10, 80);
    emit_str_x_sp_off(out, 11, 88);
    emit_str_x_sp_off(out, 12, 96);
    emit_str_x_sp_off(out, 13, 104);
    emit_str_x_sp_off(out, 14, 112);
    emit_str_x_sp_off(out, 15, 120);
    emit_str_x_sp_off(out, 16, 128);
    emit_str_x_sp_off(out, 17, 136);
    emit_str_x_sp_off(out, 30, 144);
    emit_movabs_x(out, 0, pc);
    emit_movabs_x(out, 1, ea);
    emit_movabs_x(out, 2, (uint64_t)is_store);
    emit_movabs_x(out, 3, (uint64_t)size);
    uint64_t curr_pc = (uint64_t)(*out);
    uint64_t target = (uint64_t)&dbi_log_mem_access;
    int64_t off = (int64_t)target - (int64_t)curr_pc;
    if ((off % 4) != 0 || off <= -(1LL<<27) || off >= (1LL<<27)) {
        emit_movabs_x(out, 16, target);
        emit32(out, 0xD63F0200 | (16 << 5)); // BLR x16
    } else {
        uint32_t imm26 = (uint32_t)((off >> 2) & 0x03FFFFFF);
        emit32(out, 0x94000000 | imm26);
    }
    // 恢复寄存器
    emit_ldr_x_sp_off(out, 30, 144);
    emit_ldr_x_sp_off(out, 17, 136);
    emit_ldr_x_sp_off(out, 16, 128);
    emit_ldr_x_sp_off(out, 15, 120);
    emit_ldr_x_sp_off(out, 14, 112);
    emit_ldr_x_sp_off(out, 13, 104);
    emit_ldr_x_sp_off(out, 12, 96);
    emit_ldr_x_sp_off(out, 11, 88);
    emit_ldr_x_sp_off(out, 10, 80);
    emit_ldr_x_sp_off(out, 9, 72);
    emit_ldr_x_sp_off(out, 8, 64);
    emit_ldr_x_sp_off(out, 7, 56);
    emit_ldr_x_sp_off(out, 6, 48);
    emit_ldr_x_sp_off(out, 5, 40);
    emit_ldr_x_sp_off(out, 4, 32);
    emit_ldr_x_sp_off(out, 3, 24);
    emit_ldr_x_sp_off(out, 2, 16);
    emit_ldr_x_sp_off(out, 1, 8);
    emit_ldr_x_sp_off(out, 0, 0);
    emit_add_sp_imm(out, 160);
}

void *dbi_translate_function(void *orig_fn)
{
    uint8_t *in = (uint8_t *)orig_fn;
    uint8_t *out = (uint8_t *)mc_code_cache_alloc(4096);
    if (!out) return NULL;
    uint8_t *cursor = out;
    for (int i = 0; i < 1024; i++) {
        uint32_t insn = fetch32(in);
        int is_store = 0; size_t sz = 0;
        if (is_pc_relative_or_branch(insn)) {
            // 碰到需要重定位的指令，直接跳回原始地址继续执行，避免错误
            emit_movabs_x(&cursor, 16, (uint64_t)in);
            emit_br_reg(&cursor, 16);
            break;
        } else if (is_ldr_str_unsigned(insn, &is_store, &sz)) {
            emit_call_log(&cursor, (uint64_t)in, 0, is_store, sz);
            emit32(&cursor, insn);
        } else {
            if (insn == 0xD65F03C0) { emit32(&cursor, insn); break; }
            emit32(&cursor, insn);
        }
        in += 4;
    }
    mc_clear_icache(out, cursor);
    return out;
}


